{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Introduction:\n",
    "\n",
    "\n",
    "Concrete ML (CML) supports also model compilation to the TFHE-rs format, enabling deeper interoperability between CML and TFHE-rs frameworks. This allows developers to combine CML’s intuitive, Python-based workflow with the performance and low-level control of TFHE-rs for advanced privacy-preserving applications.\n",
    "\n",
    "This notebook showcases a privacy-preserving server-side authentication flow, where access to a remote server is granted only if the user’s encrypted information satisfies a specific condition — all without ever decrypting any sensitive data throughout the process.\n",
    "\n",
    "On the server side, verification is performed using a binary classification algorithm provided by CML. The post-processing part, which includes an argmax operation on the model’s encrypted output followed by a multiplication with a random plaintext token (known only to the server), is implemented using the TFHE-rs library.\n",
    "\n",
    "This token acts as a proof of successful authentication and is returned only if the condition is met."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "No CUDA runtime is found, using CUDA_HOME='/usr/local/cuda'\n"
     ]
    }
   ],
   "source": [
    "import subprocess\n",
    "from tempfile import TemporaryDirectory\n",
    "\n",
    "import concrete_ml_extensions as fhext\n",
    "from sklearn.datasets import make_classification\n",
    "from sklearn.model_selection import train_test_split\n",
    "\n",
    "from concrete.ml.common.utils import CiphertextFormat\n",
    "from concrete.ml.deployment import FHEModelClient, FHEModelDev, FHEModelServer\n",
    "from concrete.ml.sklearn import DecisionTreeClassifier as ConcreteDecisionTreeClassifier"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Generate toy classification dataset\n",
    "x, y = make_classification(n_samples=100, class_sep=2, n_features=10, random_state=42)\n",
    "\n",
    "X_train, X_test, y_train, y_test = train_test_split(x, y, test_size=0.2, random_state=42)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<style>#sk-container-id-1 {color: black;background-color: white;}#sk-container-id-1 pre{padding: 0;}#sk-container-id-1 div.sk-toggleable {background-color: white;}#sk-container-id-1 label.sk-toggleable__label {cursor: pointer;display: block;width: 100%;margin-bottom: 0;padding: 0.3em;box-sizing: border-box;text-align: center;}#sk-container-id-1 label.sk-toggleable__label-arrow:before {content: \"▸\";float: left;margin-right: 0.25em;color: #696969;}#sk-container-id-1 label.sk-toggleable__label-arrow:hover:before {color: black;}#sk-container-id-1 div.sk-estimator:hover label.sk-toggleable__label-arrow:before {color: black;}#sk-container-id-1 div.sk-toggleable__content {max-height: 0;max-width: 0;overflow: hidden;text-align: left;background-color: #f0f8ff;}#sk-container-id-1 div.sk-toggleable__content pre {margin: 0.2em;color: black;border-radius: 0.25em;background-color: #f0f8ff;}#sk-container-id-1 input.sk-toggleable__control:checked~div.sk-toggleable__content {max-height: 200px;max-width: 100%;overflow: auto;}#sk-container-id-1 input.sk-toggleable__control:checked~label.sk-toggleable__label-arrow:before {content: \"▾\";}#sk-container-id-1 div.sk-estimator input.sk-toggleable__control:checked~label.sk-toggleable__label {background-color: #d4ebff;}#sk-container-id-1 div.sk-label input.sk-toggleable__control:checked~label.sk-toggleable__label {background-color: #d4ebff;}#sk-container-id-1 input.sk-hidden--visually {border: 0;clip: rect(1px 1px 1px 1px);clip: rect(1px, 1px, 1px, 1px);height: 1px;margin: -1px;overflow: hidden;padding: 0;position: absolute;width: 1px;}#sk-container-id-1 div.sk-estimator {font-family: monospace;background-color: #f0f8ff;border: 1px dotted black;border-radius: 0.25em;box-sizing: border-box;margin-bottom: 0.5em;}#sk-container-id-1 div.sk-estimator:hover {background-color: #d4ebff;}#sk-container-id-1 div.sk-parallel-item::after {content: \"\";width: 100%;border-bottom: 1px solid gray;flex-grow: 1;}#sk-container-id-1 div.sk-label:hover label.sk-toggleable__label {background-color: #d4ebff;}#sk-container-id-1 div.sk-serial::before {content: \"\";position: absolute;border-left: 1px solid gray;box-sizing: border-box;top: 0;bottom: 0;left: 50%;z-index: 0;}#sk-container-id-1 div.sk-serial {display: flex;flex-direction: column;align-items: center;background-color: white;padding-right: 0.2em;padding-left: 0.2em;position: relative;}#sk-container-id-1 div.sk-item {position: relative;z-index: 1;}#sk-container-id-1 div.sk-parallel {display: flex;align-items: stretch;justify-content: center;background-color: white;position: relative;}#sk-container-id-1 div.sk-item::before, #sk-container-id-1 div.sk-parallel-item::before {content: \"\";position: absolute;border-left: 1px solid gray;box-sizing: border-box;top: 0;bottom: 0;left: 50%;z-index: -1;}#sk-container-id-1 div.sk-parallel-item {display: flex;flex-direction: column;z-index: 1;position: relative;background-color: white;}#sk-container-id-1 div.sk-parallel-item:first-child::after {align-self: flex-end;width: 50%;}#sk-container-id-1 div.sk-parallel-item:last-child::after {align-self: flex-start;width: 50%;}#sk-container-id-1 div.sk-parallel-item:only-child::after {width: 0;}#sk-container-id-1 div.sk-dashed-wrapped {border: 1px dashed gray;margin: 0 0.4em 0.5em 0.4em;box-sizing: border-box;padding-bottom: 0.4em;background-color: white;}#sk-container-id-1 div.sk-label label {font-family: monospace;font-weight: bold;display: inline-block;line-height: 1.2em;}#sk-container-id-1 div.sk-label-container {text-align: center;}#sk-container-id-1 div.sk-container {/* jupyter's `normalize.less` sets `[hidden] { display: none; }` but bootstrap.min.css set `[hidden] { display: none !important; }` so we also need the `!important` here to be able to override the default hidden behavior on the sphinx rendered scikit-learn.org. See: https://github.com/scikit-learn/scikit-learn/issues/21755 */display: inline-block !important;position: relative;}#sk-container-id-1 div.sk-text-repr-fallback {display: none;}</style><div id=\"sk-container-id-1\" class=\"sk-top-container\"><div class=\"sk-text-repr-fallback\"><pre>DecisionTreeClassifier(n_bits={&#x27;op_inputs&#x27;: 8, &#x27;op_leaves&#x27;: 8})</pre><b>In a Jupyter environment, please rerun this cell to show the HTML representation or trust the notebook. <br />On GitHub, the HTML representation is unable to render, please try loading this page with nbviewer.org.</b></div><div class=\"sk-container\" hidden><div class=\"sk-item\"><div class=\"sk-estimator sk-toggleable\"><input class=\"sk-toggleable__control sk-hidden--visually\" id=\"sk-estimator-id-1\" type=\"checkbox\" checked><label for=\"sk-estimator-id-1\" class=\"sk-toggleable__label sk-toggleable__label-arrow\">DecisionTreeClassifier</label><div class=\"sk-toggleable__content\"><pre>DecisionTreeClassifier(n_bits={&#x27;op_inputs&#x27;: 8, &#x27;op_leaves&#x27;: 8})</pre></div></div></div></div></div>"
      ],
      "text/plain": [
       "DecisionTreeClassifier(n_bits={'op_inputs': 8, 'op_leaves': 8})"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# For now, only 8-bit quantization is supported for TFHE-rs interoperability\n",
    "model = ConcreteDecisionTreeClassifier(n_bits=8)\n",
    "# Train the CML model on clear data\n",
    "model.fit(X_train, y_train)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compile the model with TFHE-rs support (ciphertext_format=CiphertextFormat.TFHE_RS)\n",
    "model.compile(X_train, ciphertext_format=CiphertextFormat.TFHE_RS);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Use a temporary directory to store the Dev, client, and server models\n",
    "with TemporaryDirectory() as temp_dir:\n",
    "    # Export the model for deployment using the FHEModelDev API\n",
    "    model_dev = FHEModelDev(temp_dir, model)\n",
    "    model_dev.save()\n",
    "\n",
    "    # Loading the model on the client side using FHEModelClient API\n",
    "    fhe_model_client = FHEModelClient(path_dir=temp_dir)\n",
    "    fhe_model_client.load()\n",
    "\n",
    "    # Loading the model on the server side using FHEModelServer API\n",
    "    fhe_model_server = FHEModelServer(temp_dir)\n",
    "    fhe_model_server.load()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Client side : Generates all keys and serializes the evaluation keys required for the server\n",
    "evaluation_keys, tfhers_evaluation_keys = fhe_model_client.get_serialized_evaluation_keys(\n",
    "    include_tfhers_key=True\n",
    ")\n",
    "\n",
    "with open(\"evalkeys_tfhers.bin\", \"wb\") as fp:\n",
    "    fp.write(tfhers_evaluation_keys)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Use Case: Privacy-Preserving Authentication with Homomorphic Encryption\n",
    "\n",
    "In this scenario, the client attempts to authenticate using sensitive data (e.g., a biometric vector such as facial recognition). \n",
    "\n",
    "On the client side, the input is first quantized (to meet FHE constraints), then encrypted, and finally sent to the server.\n",
    "\n",
    "On the server side, first, a binary classification algorithm — in this case, a decision tree — is applied on the encrypted data to determine whether authentication should succeed (decision = 1) or fail (decision = 0). Then, the decision result is homomorphically multiplied by a secret authentication token known only to the server. Finally, the encrypted result is sent back to the client.\n",
    "\n",
    "On the client side, the result is decrypted using the secret key. If the authentication was successful (decision = 1), the decrypted value represents the original server token. Otherwise (decision = 0), the decrypted result is a vector of zeros, indicating authentication failure."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "🔢 Test sample #0 ==================================================\n",
      "Inference input: [[-0.07015438  2.85605469 -1.05319823  2.4898248  -0.21910053  0.32692737\n",
      "  -2.21113531  0.77086519  0.82940558  0.23561456]]\n",
      "🎯 Ground truth label: 0\n",
      "Model logits (after decryption): [1. 0.]\n",
      "🔐 Authentication token: 🛑 Access denied — token: [0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0]\n",
      "\n",
      "🔢 Test sample #1 ==================================================\n",
      "Inference input: [[-2.43430912  2.25291095  2.09885263  1.23502851 -0.24751864 -0.72713718\n",
      "   0.6206721  -1.33534436 -0.07443343  0.177701  ]]\n",
      "🎯 Ground truth label: 1\n",
      "Model logits (after decryption): [0. 1.]\n",
      "🔐 Authentication token: 🟢 Access granted — token: [-122  -15  106  -56   85   39  -28   94 -109   10   27  -51  -33   60\n",
      "  -25   11] \n",
      "\n",
      "🔢 Test sample #2 ==================================================\n",
      "Inference input: [[-3.37679987  3.30450019  2.83991037  1.87087893 -0.37482081 -0.2403254\n",
      "   0.44426331  1.1593298   0.71095997 -0.36096617]]\n",
      "🎯 Ground truth label: 1\n",
      "Model logits (after decryption): [0. 1.]\n",
      "🔐 Authentication token: 🟢 Access granted — token: [ 117  -34 -113   75  -19   55 -125  116   18   94   35  100   -5  -87\n",
      "   67  -95] \n",
      "\n",
      "🔢 Test sample #3 ==================================================\n",
      "Inference input: [[-1.97994267  1.53479393  1.82584805  0.74282242 -0.47874862  0.22213377\n",
      "  -0.8946073  -0.43973106  1.25575613 -0.18687164]]\n",
      "🎯 Ground truth label: 1\n",
      "Model logits (after decryption): [0. 1.]\n",
      "🔐 Authentication token: 🟢 Access granted — token: [ -20   45  103    3    7  -47  -58   -4 -111   82  -46   52  -56  -10\n",
      "  -31  -30] \n",
      "\n",
      "🔢 Test sample #4 ==================================================\n",
      "Inference input: [[-1.19229599 -1.71352532  2.15199146 -1.87205691 -0.59939265 -0.83972184\n",
      "  -0.52575502  0.15039379 -2.12389572 -0.75913266]]\n",
      "🎯 Ground truth label: 1\n",
      "Model logits (after decryption): [0. 1.]\n",
      "🔐 Authentication token: 🟢 Access granted — token: [  -9 -112  102   71 -122   16  -54    5  -49   94  -29  -34  -52  -73\n",
      "  -35  125] \n",
      "\n",
      "🔢 Test sample #5 ==================================================\n",
      "Inference input: [[ 2.48305377 -2.20592943 -2.17762637 -1.17878084  0.23204994 -0.47103831\n",
      "  -1.40746377 -0.21344715 -1.44808434 -0.71844422]]\n",
      "🎯 Ground truth label: 0\n",
      "Model logits (after decryption): [1. 0.]\n",
      "🔐 Authentication token: 🛑 Access denied — token: [0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0]\n",
      "\n",
      "🔢 Test sample #6 ==================================================\n",
      "Inference input: [[ 2.85855887 -2.43534875 -2.54851014 -1.26544183  0.40171172 -0.10876015\n",
      "  -0.40122047  0.0125924   0.69014399  0.22409248]]\n",
      "🎯 Ground truth label: 0\n",
      "Model logits (after decryption): [1. 0.]\n",
      "🔐 Authentication token: 🛑 Access denied — token: [0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0]\n",
      "\n",
      "🔢 Test sample #7 ==================================================\n",
      "Inference input: [[-2.02146424  2.02703224  1.68057323  1.16292288  0.70030988 -0.85835778\n",
      "   0.12200981 -0.0960599  -0.57563783  2.56008454]]\n",
      "🎯 Ground truth label: 1\n",
      "Model logits (after decryption): [0. 1.]\n",
      "🔐 Authentication token: 🟢 Access granted — token: [-117   97  -78   21    7 -104  106  -89 -104   77  -55   37   64   80\n",
      "  -41  -96] \n",
      "\n",
      "🔢 Test sample #8 ==================================================\n",
      "Inference input: [[-1.96346177  1.88629356  1.66530077  1.05694552 -1.12905177  0.72167206\n",
      "   0.48937456  0.71299843 -0.52452027 -1.22212781]]\n",
      "🎯 Ground truth label: 1\n",
      "Model logits (after decryption): [0. 1.]\n",
      "🔐 Authentication token: 🟢 Access granted — token: [  85  -83   64 -105 -100 -116  -19  -71   31  -70  108   10   33   46\n",
      " -118   93] \n",
      "\n",
      "🔢 Test sample #9 ==================================================\n",
      "Inference input: [[-2.27657182  2.35970566  1.86199147  1.37726866  0.60600995 -1.55662917\n",
      "   1.75479418  1.69645637 -1.28042935 -2.08192941]]\n",
      "🎯 Ground truth label: 1\n",
      "Model logits (after decryption): [0. 1.]\n",
      "🔐 Authentication token: 🟢 Access granted — token: [ -50  -97  -99 -115   -1  -10  -92   46 -115 -111 -115  127  -24   38\n",
      "  125   92] \n",
      "\n"
     ]
    }
   ],
   "source": [
    "for DATA_INDEX in range(10):\n",
    "\n",
    "    print(f\"🔢 Test sample #{DATA_INDEX}\", \"=\" * 50)\n",
    "    print(\"Inference input:\", X_test[[DATA_INDEX]])\n",
    "    print(f\"🎯 Ground truth label: {y_test[DATA_INDEX]}\")\n",
    "\n",
    "    # Client side : Quantize, encrypt and serialize the data\n",
    "    q_x_encrypted_serialized = fhe_model_client.quantize_encrypt_serialize(X_test[[DATA_INDEX]])\n",
    "\n",
    "    # Server side: Run the model over encrypted data\n",
    "    q_y_pred_encrypted_serialized = fhe_model_server.run(q_x_encrypted_serialized, evaluation_keys)\n",
    "\n",
    "    # Client side : Decrypt, de-quantize and post-process the result\n",
    "    y_pred = fhe_model_client.deserialize_decrypt_dequantize(q_y_pred_encrypted_serialized[0])\n",
    "\n",
    "    print(\"Model logits (after decryption):\", y_pred.reshape((-1,)))\n",
    "\n",
    "    with open(\"prediction_non_preprocessed.bin\", \"wb\") as f:\n",
    "        f.write(q_y_pred_encrypted_serialized[0])\n",
    "\n",
    "    # ***************************************\n",
    "    # NOW RUN cargo run --release\n",
    "    # **************************************\n",
    "\n",
    "    subprocess.run(\n",
    "        [\n",
    "            \"cargo\",\n",
    "            \"run\",\n",
    "            \"--release\",\n",
    "            \"--\",  # Required to pass args to the Rust binary\n",
    "            \"evalkeys_tfhers.bin\",\n",
    "            \"prediction_non_preprocessed.bin\",\n",
    "            \"tfhers_sign_result.bin\",\n",
    "        ],\n",
    "        capture_output=True,\n",
    "        text=True,\n",
    "        check=True,\n",
    "    )\n",
    "\n",
    "    with open(\"tfhers_sign_result.bin\", \"rb\") as f:\n",
    "        tfhers_postproc_bytes = f.read()\n",
    "\n",
    "    results = fhext.decrypt_radix(\n",
    "        blob=tfhers_postproc_bytes,\n",
    "        shape=(1, 16),\n",
    "        bitwidth=8,\n",
    "        is_signed=True,\n",
    "        secret_key=fhe_model_client.get_tfhers_secret_key(),\n",
    "    )\n",
    "\n",
    "    token = results.reshape((-1,))\n",
    "\n",
    "    if all(t == 0 for t in token):\n",
    "        print(f\"🔐 Authentication token: 🛑 Access denied — token: {token}\\n\")\n",
    "    else:\n",
    "        print(f\"🔐 Authentication token: 🟢 Access granted — token: {token} \\n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Conclusion\n",
    "\n",
    "In this notebook, we showed how to combine Concrete ML (CML) and TFHE-rs for end-to-end encrypted processing. After training and compiling a CML decision tree, we performed an encrypted inference and then applied additional post-processing in TFHE-rs, all while preserving data privacy.\n",
    "\n",
    "This approach highlights the power of interoperability between CML and TFHE-rs, and paves the way for more advanced privacy-preserving use cases."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Disclaimer**: This use case assumes an honest-but-curious server and does not represent a complete secure authentication protocol. Further measures are required to protect against malicious servers."
   ]
  }
 ],
 "metadata": {
  "execution": {
   "timeout": 10800
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
